import UIKit

public class SyntaxHighlighter {
    private let theme: SyntaxTheme
    
    public init(theme: SyntaxTheme) {
        self.theme = theme
    }
    
    public func highlight(_ code: String, language: String?) -> NSAttributedString {
        let attributedString = NSMutableAttributedString(string: code, attributes: [
            .foregroundColor: theme.plainText.color,
            .font: theme.plainText.font ?? .monospacedSystemFont(ofSize: 14, weight: .regular)
        ])
        
        guard let language = language?.lowercased() else {
            return attributedString
        }
        
        switch language {
        case "swift":
            highlightSwift(in: attributedString)
        case "python":
            highlightPython(in: attributedString)
        case "javascript", "js":
            highlightJavaScript(in: attributedString)
        case "html":
            highlightHTML(in: attributedString)
        case "css":
            highlightCSS(in: attributedString)
        case "json":
            highlightJSON(in: attributedString)
        case "markdown", "md":
            highlightMarkdown(in: attributedString)
        case "xml":
            highlightXML(in: attributedString)
        case "sql":
            highlightSQL(in: attributedString)
        case "yaml", "yml":
            highlightYAML(in: attributedString)
        case "bash", "sh":
            highlightBash(in: attributedString)
        default:
            highlightSwift(in: attributedString)
            break
        }
        
        return attributedString
    }
    
    private func highlightPattern(_ pattern: String, in code: NSMutableAttributedString, tokenType: SyntaxTheme.TokenType) {
        guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { return }
        let range = NSRange(location: 0, length: code.length)
        let matches = regex.matches(in: code.string, options: [], range: range)
        
        for match in matches.reversed() {
            var attributes: [NSAttributedString.Key: Any] = [.foregroundColor: tokenType.color]
            if let font = tokenType.font {
                attributes[.font] = font
            }
            if let backgroundColor = tokenType.backgroundColor {
                attributes[.backgroundColor] = backgroundColor
            }
            code.setAttributes(attributes, range: match.range)
        }
    }
    
    private func highlightSwift(in code: NSMutableAttributedString) {
        // Keywords
        highlightPattern("\\b(class|struct|enum|protocol|extension|func|var|let|if|else|guard|switch|case|default|for|while|do|try|catch|throw|throws|rethrows|return|break|continue|fallthrough|defer|where|in|init|deinit|self|super|override|subscript|convenience|weak|unowned|required|lazy|get|set|willSet|didSet|mutating|nonmutating|prefix|infix|postfix|precedencegroup|associativity|left|right|none|operator|import|typealias|as|is|any|some|await|async|actor|public|internal|private|fileprivate|open|final|static|dynamic|optional|required|convenience|indirect|associatedtype|precedencegroup)\\b", in: code, tokenType: theme.keyword)
        
        // Types
        highlightPattern("\\b(String|Int|Double|Float|Bool|Character|Array|Dictionary|Set|Void|Any|AnyObject|Error|Optional)\\b", in: code, tokenType: theme.type)
        
        // Numbers
        highlightPattern("\\b\\d+(\\.\\d+)?\\b", in: code, tokenType: theme.number)
        
        // String literals
        highlightPattern("\"[^\"\\\\]*(?:\\\\.[^\"\\\\]*)*\"", in: code, tokenType: theme.string)
        
        // Comments
        highlightPattern("//.*$|/\\*[\\s\\S]*?\\*/", in: code, tokenType: theme.comment)
        
        // Custom types (capitalized words)
        highlightPattern("\\b[A-Z][a-zA-Z0-9_]*\\b", in: code, tokenType: theme.customType)
        
        // Attributes
        highlightPattern("@\\w+\\b", in: code, tokenType: theme.attribute)
    }
    
    private func highlightPython(in code: NSMutableAttributedString) {
        // Keywords
        highlightPattern("\\b(and|as|assert|break|class|continue|def|del|elif|else|except|finally|for|from|global|if|import|in|is|lambda|nonlocal|not|or|pass|raise|return|try|while|with|yield)\\b", in: code, tokenType: theme.keyword)
        
        // Built-in types
        highlightPattern("\\b(bool|bytearray|bytes|dict|float|frozenset|int|list|object|set|str|tuple|type)\\b", in: code, tokenType: theme.type)
        
        // Numbers
        highlightPattern("\\b\\d+(\\.\\d+)?\\b", in: code, tokenType: theme.number)
        
        // String literals
        highlightPattern("(['\"])[^'\"\\\\]*(?:\\\\.[^'\"\\\\]*)*\\1", in: code, tokenType: theme.string)
        
        // Comments
        highlightPattern("#.*$", in: code, tokenType: theme.comment)
        
        // Decorators
        highlightPattern("@\\w+\\b", in: code, tokenType: theme.attribute)
    }
    
    private func highlightJavaScript(in code: NSMutableAttributedString) {
        // Keywords
        highlightPattern("\\b(break|case|catch|class|const|continue|debugger|default|delete|do|else|export|extends|finally|for|function|if|import|in|instanceof|new|return|super|switch|this|throw|try|typeof|var|void|while|with|yield|let|static|enum|await|implements|package|protected|interface|private|public|arguments|eval|null|undefined|true|false)\\b", in: code, tokenType: theme.keyword)
        
        // Built-in types and objects
        highlightPattern("\\b(Array|Boolean|Date|Error|Function|JSON|Math|Number|Object|RegExp|String|Promise|Map|Set|Symbol|WeakMap|WeakSet)\\b", in: code, tokenType: theme.type)
        
        // Numbers
        highlightPattern("\\b\\d+(\\.\\d+)?\\b", in: code, tokenType: theme.number)
        
        // String literals
        highlightPattern("(['\"`])[^'\"`\\\\]*(?:\\\\.[^'\"`\\\\]*)*\\1", in: code, tokenType: theme.string)
        
        // Comments
        highlightPattern("//.*$|/\\*[\\s\\S]*?\\*/", in: code, tokenType: theme.comment)
        
        // Regular expressions
        highlightPattern("/[^/\\\\]*(?:\\\\.[^/\\\\]*)*/[gimuy]*", in: code, tokenType: theme.regex)
    }
    
    private func highlightHTML(in code: NSMutableAttributedString) {
        // Tags
        highlightPattern("</?[^>]+>", in: code, tokenType: theme.keyword)
        
        // Attributes
        highlightPattern("\\b\\w+(?=\\s*=\\s*['\"])", in: code, tokenType: theme.attribute)
        
        // String literals
        highlightPattern("\"[^\"\\\\]*(?:\\\\.[^\"\\\\]*)*\"|'[^'\\\\]*(?:\\\\.[^'\\\\]*)*'", in: code, tokenType: theme.string)
        
        // Comments
        highlightPattern("<!--[\\s\\S]*?-->", in: code, tokenType: theme.comment)
        
        // DOCTYPE
        highlightPattern("<!DOCTYPE[^>]*>", in: code, tokenType: theme.preprocessor)
    }
    
    private func highlightCSS(in code: NSMutableAttributedString) {
        // Selectors
        highlightPattern("[^{}\\s][^{}]*(?=\\s*\\{)", in: code, tokenType: theme.keyword)
        
        // Properties
        highlightPattern("\\b[a-z-]+(?=\\s*:)", in: code, tokenType: theme.attribute)
        
        // Values
        highlightPattern(":\\s*[^;\\n]+", in: code, tokenType: theme.string)
        
        // Comments
        highlightPattern("/\\*[\\s\\S]*?\\*/", in: code, tokenType: theme.comment)
        
        // Units
        highlightPattern("\\b\\d+(\\.\\d+)?(px|em|rem|%|vh|vw|vmin|vmax|ch|ex|cm|mm|in|pt|pc|fr)\\b", in: code, tokenType: theme.number)
        
        // Colors
        highlightPattern("#([A-Fa-f0-9]{6}|[A-Fa-f0-9]{3})\\b", in: code, tokenType: theme.string)
    }
    
    private func highlightJSON(in code: NSMutableAttributedString) {
        // String literals
        highlightPattern("\"[^\"\\\\]*(?:\\\\.[^\"\\\\]*)*\"(?=\\s*:)?", in: code, tokenType: theme.attribute)
        highlightPattern("(?<=:)\\s*\"[^\"\\\\]*(?:\\\\.[^\"\\\\]*)*\"", in: code, tokenType: theme.string)
        
        // Numbers
        highlightPattern("\\b-?\\d+(\\.\\d+)?([eE][+-]?\\d+)?\\b", in: code, tokenType: theme.number)
        
        // Keywords
        highlightPattern("\\b(true|false|null)\\b", in: code, tokenType: theme.keyword)
    }
    
    private func highlightMarkdown(in code: NSMutableAttributedString) {
        // Headers
        highlightPattern("^#{1,6}\\s.*$", in: code, tokenType: theme.keyword)
        
        // Emphasis
        highlightPattern("\\*\\*[^\\*\\n]+\\*\\*|__[^_\\n]+__", in: code, tokenType: theme.type)
        highlightPattern("\\*[^\\*\\n]+\\*|_[^_\\n]+_", in: code, tokenType: theme.identifier)
        
        // Links
        highlightPattern("\\[([^\\]]+)\\]\\([^\\)]+\\)", in: code, tokenType: theme.url)
        
        // Code blocks
        highlightPattern("`[^`]+`|```[\\s\\S]*?```", in: code, tokenType: theme.string)
        
        // Lists
        highlightPattern("^[\\*\\-\\+]\\s.*$", in: code, tokenType: theme.attribute)
        highlightPattern("^\\d+\\.\\s.*$", in: code, tokenType: theme.number)
    }
    
    private func highlightXML(in code: NSMutableAttributedString) {
        // Tags
        highlightPattern("</?[^>]+>", in: code, tokenType: theme.keyword)
        
        // Attributes
        highlightPattern("\\b\\w+(?=\\s*=\\s*['\"])", in: code, tokenType: theme.attribute)
        
        // String literals
        highlightPattern("\"[^\"\\\\]*(?:\\\\.[^\"\\\\]*)*\"|'[^'\\\\]*(?:\\\\.[^'\\\\]*)*'", in: code, tokenType: theme.string)
        
        // Comments
        highlightPattern("<!--[\\s\\S]*?-->", in: code, tokenType: theme.comment)
        
        // CDATA
        highlightPattern("<!\\[CDATA\\[[\\s\\S]*?\\]\\]>", in: code, tokenType: theme.preprocessor)
    }
    
    private func highlightSQL(in code: NSMutableAttributedString) {
        // Keywords
        highlightPattern("\\b(SELECT|INSERT|UPDATE|DELETE|FROM|WHERE|GROUP BY|ORDER BY|HAVING|JOIN|LEFT|RIGHT|INNER|OUTER|ON|AS|AND|OR|NOT|IN|BETWEEN|LIKE|IS|NULL|CREATE|ALTER|DROP|TABLE|INDEX|VIEW|TRIGGER|PROCEDURE|FUNCTION|DATABASE|GRANT|REVOKE|COMMIT|ROLLBACK|BEGIN|END|DECLARE|SET)\\b", in: code, tokenType: theme.keyword)
        
        // Functions
        highlightPattern("\\b(COUNT|SUM|AVG|MIN|MAX|FIRST|LAST|UCASE|LCASE|MID|LEN|ROUND|NOW|FORMAT)\\b", in: code, tokenType: theme.type)
        
        // Numbers
        highlightPattern("\\b\\d+(\\.\\d+)?\\b", in: code, tokenType: theme.number)
        
        // String literals
        highlightPattern("'[^'\\\\]*(?:\\\\.[^'\\\\]*)*'", in: code, tokenType: theme.string)
        
        // Comments
        highlightPattern("--.*$|/\\*[\\s\\S]*?\\*/", in: code, tokenType: theme.comment)
    }
    
    private func highlightYAML(in code: NSMutableAttributedString) {
        // Keys
        highlightPattern("^\\s*[^:]+(?=:)", in: code, tokenType: theme.attribute)
        
        // String literals
        highlightPattern(":\\s*(['\"])[^'\"]*\\1", in: code, tokenType: theme.string)
        
        // Numbers
        highlightPattern(":\\s*-?\\d+(\\.\\d+)?\\b", in: code, tokenType: theme.number)
        
        // Comments
        highlightPattern("#.*$", in: code, tokenType: theme.comment)
        
        // Keywords
        highlightPattern("\\b(true|false|null|~)\\b", in: code, tokenType: theme.keyword)
    }
    
    private func highlightBash(in code: NSMutableAttributedString) {
        // Keywords
        highlightPattern("\\b(if|then|else|elif|fi|case|esac|for|while|until|do|done|in|function|time|select)\\b", in: code, tokenType: theme.keyword)
        
        // Built-in commands
        highlightPattern("\\b(echo|exit|return|set|unset|export|alias|unalias|cd|pwd|ls|cp|mv|rm|mkdir|rmdir|touch|cat|less|more|head|tail|grep|sed|awk|find|sort|uniq|wc|date|curl|wget|ssh|scp|tar|gzip|gunzip|zip|unzip)\\b", in: code, tokenType: theme.type)
        
        // Variables
        highlightPattern("\\$\\w+|\\$\\{[^}]+\\}", in: code, tokenType: theme.identifier)
        
        // String literals
        highlightPattern("\"[^\"\\\\]*(?:\\\\.[^\"\\\\]*)*\"|'[^'\\\\]*(?:\\\\.[^'\\\\]*)*'", in: code, tokenType: theme.string)
        
        // Comments
        highlightPattern("#.*$", in: code, tokenType: theme.comment)
        
        // Options
        highlightPattern("\\s-\\w+\\b|\\s--[\\w-]+\\b", in: code, tokenType: theme.attribute)
    }
}
